#![feature(trait_alias)]

extern crate chatbox;

mod grammarfile;

#[cfg(not(test))]
use grammarfile::{further_validate_grammar, parse_grammar, syntactically_validate_grammar, GetOk};

mod util;
#[cfg(not(test))]
use util::*;

mod crypto;
#[cfg(not(test))]
use crypto::KEY_SIZE;

mod encode;
#[cfg(not(test))]
use encode::*;

mod decode;
#[cfg(not(test))]
use decode::*;

mod packetizer;
#[cfg(not(test))]
use packetizer::*;

mod testlib;
#[cfg(test)]
use testlib::*;

#[macro_use]
extern crate lazy_static;

#[cfg(test)]
#[macro_use]
extern crate matches;

#[cfg(not(test))]
extern crate clap;
#[cfg(not(test))]
use clap::{Arg, ArgGroup, Command};

#[cfg(not(test))]
extern crate serde_yaml;

#[cfg(not(test))]
use std::fs::File;
#[cfg(not(test))]
use std::io::prelude::*;
#[cfg(not(test))]
use std::io::BufReader;

// Not regex::bytes::Regex: used for reading YAML which is always unicode
#[cfg(not(test))]
use regex::Regex;
extern crate base64;

// Inspired from, and similar to the method described in: http://www.nic.funet.fi/pub/crypt/old/mimic/mimic.text
// See also: https://news.ycombinator.com/item?id=8721475

use chatbox::*;

#[cfg(not(test))]
fn main() -> Result<(), anyhow::Error> {
    let matches = Command::new("chatbox")
        .version("0.0.9-alpha")
        .author("fnordomat <GPG:46D46D1246803312401472B5A7427E237B7908CA>")
        .about("Communication using ordered stochastic context-free grammars.")
        .arg(
            Arg::new("grammar")
                .short('g')
                .long("grammar")
                .required(true)
                .num_args(1)
                .help("Load grammar from file"),
        )
        .arg(
            Arg::new("encode")
                .short('e')
                .long("encode")
                .num_args(0)
                .help("Act as an Encoder. Note: when -q [crypto.yaml] is also given, will cut the input into packets and encrypt + encode each."),
        )
        .arg(
            Arg::new("decode")
                .short('d')
                .long("decode")
                .num_args(0)
                .help("Act as a Decoder (Arbitrary Unambiguous Grammar). Externally delimited, expects format ([LENGTH(HEX)]\n[MSG])*"),
        )
        .arg(
            Arg::new("prefixfree")
                .short('p')
                .long("prefixfree")
                // why is there no requires_one? .requires_one(&["encode", "decode"])
                .num_args(0)
                .help("Prefix-free language mode. Code sentences are self-delimiting: pipe a stream, no length-prefix output (with --encode) or expected (with --decode)"),
        )
        .arg(
            Arg::new("input")
                .short('i')
                .long("input")
                .num_args(1)
                .help("Input file"))
        .arg(
            Arg::new("validate")
                .short('v')
                .long("validate")
                .num_args(0)
                .help("Just validate the grammar"))
        .arg(
            Arg::new("crypto")
                .short('q')
                .long("crypto")
                .num_args(1)
                .help("Use encrypted and authenticated packet format, using pre-shared secret from given file"))
        .group(ArgGroup::new("mode")
               .args(&["encode", "decode", "validate"])
               .required(true)).get_matches();

    let prefixfree = matches.get_flag("prefixfree");

    let grammar_file_path = std::path::PathBuf::from(matches.get_one::<String>("grammar").unwrap());

    let grammar_file = File::open(grammar_file_path)?;
    let mut buf_reader = BufReader::new(grammar_file);
    let mut contents = String::new();
    buf_reader.read_to_string(&mut contents)?;

    // contents holds a unicode string, but our format can represent arbitrary
    // bytestrings using hex codes in angle brackets, for example <C0 FF EE>

    let grammar = parse_grammar(&contents)?;

    use std::io;

    // Key used to secure (encrypt, MAC) the channel
    let maybe_key = match matches.get_one::<String>("crypto") {
        Some(crypto_file_path) => {
            let crypto_file = File::open(crypto_file_path)?;
            let mut buf_reader = BufReader::new(crypto_file);
            let mut contents = String::new();
            buf_reader.read_to_string(&mut contents)?;

            let res: serde_yaml::Value = serde_yaml::from_str(&contents)?;
            let psk = res.get_ok("crypto")?.get_ok("psk")?.clone();
            let psk: String = psk
                .as_str()
                .ok_or_else(|| anyhow::Error::msg("yaml error: crypto/psk is not a string value"))?
                .to_string();

            lazy_static! {
                static ref STRING_ENCODED_VALUE: Regex =
                    Regex::new(r#"base64:(?P<base64>[A-Za-z0-9+/=]+)"#).unwrap();
            }

            match &STRING_ENCODED_VALUE
                .captures_iter(&psk[..])
                .collect::<Vec<_>>()[..]
            {
                [b64] => {
                    let psk = base64::decode(
                        b64.name("base64")
                            .ok_or_else(|| {
                                anyhow::Error::msg(
                                    "error: unable to find base64encoded value of psk",
                                )
                            })?
                            .as_str(),
                    )?;
                    Some(psk)
                }
                _ => None,
            }
        }
        None => None,
    };

    if matches.get_flag("validate") {
        let val = syntactically_validate_grammar(&grammar);
        println!("{}", val);
        let vals = further_validate_grammar(&grammar);
        println!("Further issues: {:?}", vals);
    } else if matches.get_flag("encode") {
        let reader: Box<dyn io::Read> = match matches.get_one::<String>("input") {
            Some(filename) => {
                let file = File::open(filename)?;
                Box::new(io::BufReader::new(file))
            }
            None => Box::new(io::BufReader::new(io::stdin())),
        };

        let iter = reader
            .bytes()
            .into_iter()
            .take_while(|x| x.is_ok())
            .map(|x| x.unwrap());

        use itertools::Itertools;
        use rand::Rng;

        match maybe_key {
            Some(key) => {
                // Arbitrary, not too large
                let chunk_size = 32;
                let mut stream_id: [u8; 8] = Default::default();
                let mut rng = rand::thread_rng();
                rng.fill(&mut stream_id[..]);
                let chunks = iter.into_iter().chunks(chunk_size);
                let mut it = chunks.into_iter().map(|chunk| chunk.collect());
                let mut packetizer = Packetizer::with_live_random(
                    &key[0..*KEY_SIZE],
                    &mut it,
                    // chunk_size,
                    &stream_id,
                );

                loop {
                    let next = packetizer.next().clone();
                    match next {
                        Some(output) => {
                            let encode = Encoder::with_live_random(&grammar, &output[..]);
                            for s in encode {
                                if prefixfree {
                                    io::stdout().lock().write_all(&s[..])?;
                                } else {
                                    println!("{}", s.len());
                                    io::stdout().lock().write_all(&s[..])?;
                                    println!();
                                }
                            }
                        }
                        None => {
                            break;
                        }
                    }
                }
            }
            // Simply encode the input (one shot)
            None => {
                let vec: Vec<u8> = iter.collect();
                let encode = Encoder::with_live_random(&grammar, &vec[..]);
                for s in encode {
                    if prefixfree {
                        io::stdout().lock().write_all(&s[..])?;
                    } else {
                        println!("{}", s.len());
                        io::stdout().lock().write_all(&s[..])?;
                        println!();
                    }
                }
            }
        }

    // TODO -> what's also still missing: re-assemble each (chunked) packet
    // (according to conv. id and serial number)
    } else if matches.get_flag("decode") {
        if matches.get_flag("prefixfree") {
            let reader: Box<dyn io::Read> = match matches.get_one::<String>("input") {
                Some(filename) => {
                    let file = File::open(filename)?;
                    Box::new(io::BufReader::new(file))
                }
                None => Box::new(io::BufReader::new(io::stdin())),
            };

            let mut input_iter = reader
                .bytes()
                .into_iter()
                .take_while(|x| x.is_ok())
                .map(|x| x.unwrap());

            match maybe_key {
                Some(key) => {
                    let p_dec = PrefixfreeDecoder::new(&grammar, &mut input_iter);

                    let depacketizer = DecodingLoop::new(
                        Box::new(p_dec),
                        Box::new(|fragments| {
                            trial_decode_header_from_fragments(&key[0..*KEY_SIZE], fragments)
                        }),
                        Box::new(|size, fragments| {
                            trial_decode_from_fragments(&key[0..*KEY_SIZE], size, fragments)
                        }),
                    );

                    for packet in depacketizer {
                        // debug
                        // println!("{:?}", packet);
                        io::stdout().lock().write_all(&packet.payload)?;
                    }
                }
                None => {
                    let mut dec = PrefixfreeDecoder::new(&grammar, &mut input_iter);
                    let msg = dec.decode().unwrap();
                    std::io::stdout().write_all(&msg)?;
                    std::io::stdout().flush()?;
                }
            }
        // not prefixfree case
        } else {
            let reader: Box<dyn io::Read> = match matches.get_one::<String>("input") {
                Some(filename) => {
                    let file = File::open(filename)?;
                    Box::new(io::BufReader::new(file))
                }
                None => Box::new(io::BufReader::new(io::stdin())),
            };

            let mut iter = reader
                .bytes()
                .into_iter()
                .take_while(|x| x.is_ok())
                .map(|x| x.unwrap());
            let mut ld_iter = LengthDelimitedInputIterator::new(&mut iter);

            // TODO: handle empty, incomplete or wrong input gracefully

            match maybe_key {
                Some(key) => {
                    let bounds_iterator = ld_iter.filter_map(|w| {
                        let inputs = vec![w];
                        let mut iter = inputs.into_iter();
                        let mut dec = Decoder::new(&grammar, &mut iter);
                        match dec.decode() {
                            Ok(ok) => Some(ok),
                            _ => None,
                        }
                    });

                    let depacketizer = DecodingLoop::new(
                        Box::new(bounds_iterator),
                        Box::new(|fragments| {
                            trial_decode_header_from_fragments(&key[0..*KEY_SIZE], fragments)
                        }),
                        Box::new(|size, fragments| {
                            trial_decode_from_fragments(&key[0..*KEY_SIZE], size, fragments)
                        }),
                    );

                    for packet in depacketizer {
                        // debug
                        // println!("{:?}", packet);
                        io::stdout().lock().write_all(&packet.payload)?;
                    }
                }
                // just decode all at once
                None => {
                    let mut dec = Decoder::new(&grammar, &mut ld_iter);
                    let res = dec.decode();

                    match res {
                        Ok(result) => {
                            let s = binaryfraction_between_0_and_1_to_bytes(&result.l);
                            io::stdout().lock().write_all(&s[..])?;
                        }
                        Err(e) => {
                            eprintln!("{}", e);
                        }
                    }
                }
            }
        }
    }

    Ok(())
}
